// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//                            avtXGCFileFormat.C                           //
// ************************************************************************* //

#include <avtXGCFileFormat.h>

#include <avtDatabaseMetaData.h>

#include <InvalidVariableException.h>
#include <InvalidDBTypeException.h>
#include <InvalidFilesException.h>

#include <DebugStream.h>

#include <vtkFloatArray.h>
#include <vtkGenericCell.h>
#include <vtkUnstructuredGrid.h>

// ****************************************************************************
//  Method: avtXGCFileFormat constructor
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Jun 7 11:27:14 PDT 2016
//
// ****************************************************************************

avtXGCFileFormat::avtXGCFileFormat(const char *filename)
    : avtSTSDFileFormat(filename)
{
    numNodes = 0;
    numPhi = 0;
    initialized = false;

    // Make sure that the file is in fact GTC.
    Initialize();
}


// ****************************************************************************
//  Method: avtXGCFileFormat::FreeUpResources
//
//  Purpose:
//      When VisIt is done focusing on a particular timestep, it asks that
//      timestep to free up any resources (memory, file descriptors) that
//      it has associated with it.  This method is the mechanism for doing
//      that.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Jun 7 11:27:14 PDT 2016
//
// ****************************************************************************

void
avtXGCFileFormat::FreeUpResources(void)
{
}


// ****************************************************************************
//  Method: avtXGCFileFormat::PopulateDatabaseMetaData
//
//  Purpose:
//      This database meta-data object is like a table of contents for the
//      file.  By populating it, you are telling the rest of VisIt what
//      information it can request from you.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Jun 7 11:27:14 PDT 2016
//
// ****************************************************************************

void
avtXGCFileFormat::PopulateDatabaseMetaData(avtDatabaseMetaData *md)
{
    AddMeshToMetaData(md, "mesh", AVT_UNSTRUCTURED_MESH, NULL, 1, 0, 3, 3);
    AddMeshToMetaData(md, "mesh2D", AVT_UNSTRUCTURED_MESH, NULL, 1, 0, 2, 2);

    //Pull out the variables.
    hid_t vfile = H5Fopen(GetFilename(), H5F_ACC_RDONLY, H5P_DEFAULT);
    hid_t grp = H5Gopen(vfile, "/", H5P_DEFAULT);
    hsize_t num;
    H5Gget_num_objs(grp, &num);

    for (int i = 0; i < num; i++)
    {
        char nm[1024];
        H5Gget_objname_by_idx(grp,i,nm,1024);
        int otype =  H5Gget_objtype_by_idx(grp, i);
        if (otype == H5G_DATASET)
        {
            hsize_t dims[2] = {0,0};
            hid_t DS = H5Dopen(grp, nm, H5P_DEFAULT);
            hid_t vS = H5Dget_space(DS);
            int n = H5Sget_simple_extent_dims(vS, dims, NULL);

            if (n == 1 && dims[0] == numNodes)
                AddScalarVarToMetaData(md, nm, "mesh2D", AVT_NODECENT);
            else if (n == 2 && dims[0] == numNodes && dims[1] == numPhi)
                AddScalarVarToMetaData(md, nm, "mesh", AVT_NODECENT);

            H5Sclose(vS);
            H5Dclose(DS);
        }
    }

    H5Gclose(grp);
    H5Fclose(vfile);
}


// ****************************************************************************
//  Method: avtXGCFileFormat::GetMesh
//
//  Purpose:
//      Gets the mesh associated with this file.  The mesh is returned as a
//      derived type of vtkDataSet (ie vtkRectilinearGrid, vtkStructuredGrid,
//      vtkUnstructuredGrid, etc).
//
//  Arguments:
//      meshname    The name of the mesh of interest.  This can be ignored if
//                  there is only one mesh.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Jun 7 11:27:14 PDT 2016
//
// ****************************************************************************

vtkDataSet *
avtXGCFileFormat::GetMesh(const char *meshname)
{
    std::string meshFile = GetFilename();
    std::string::size_type i0 = meshFile.rfind("xgc.");
    std::string::size_type i1 = meshFile.rfind(".h5");
    
    if (i0 != std::string::npos && i1 != std::string::npos)
        meshFile = meshFile.substr(0,i0+4) + "mesh.h5";
    else
        EXCEPTION1(InvalidVariableException, "Invalid mesh filename");

    hsize_t dims[3];
    hid_t file;
    
    if ((file = H5Fopen(meshFile.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT)) < 0)
        EXCEPTION1(InvalidFilesException, meshFile);

    //Read connectivity.
    hid_t cellG = H5Gopen(file, "/cell_set[0]", H5P_DEFAULT);
    hid_t cellD = H5Dopen(cellG, "node_connect_list", H5P_DEFAULT);
    hid_t cellS = H5Dget_space(cellD);
    H5Sget_simple_extent_dims(cellS, dims, NULL);
    int *conn = new int[dims[0]*dims[1]];
    int numTri = dims[0];
    H5Dread(cellD, H5T_NATIVE_INT, H5S_ALL, cellS, H5P_DEFAULT, conn);
    H5Sclose(cellS);
    H5Dclose(cellD);
    H5Gclose(cellG);

    //Read RZ points.
    hid_t coordG = H5Gopen(file, "/coordinates", H5P_DEFAULT);
    hid_t coordD = H5Dopen(coordG, "values", H5P_DEFAULT);
    hid_t coordS = H5Dget_space(coordD);
    H5Sget_simple_extent_dims(coordS, dims, NULL);
    float *RZ = new float[dims[0]*dims[1]];

    H5Dread(coordD, H5T_NATIVE_FLOAT, H5S_ALL, coordS, H5P_DEFAULT, RZ);

    H5Sclose(coordS);
    H5Dclose(coordD);
    H5Gclose(coordG);
    
    vtkUnstructuredGrid  *grid = vtkUnstructuredGrid::New();
    vtkPoints *points = vtkPoints::New();
    grid->SetPoints(points);

    //2D mesh
    if (!strcmp(meshname, "mesh2D"))
    {
        points->SetNumberOfPoints(numNodes);
        int p = 0;
        for (int i = 0; i < numNodes; i++, p+=2)
            points->SetPoint(i, RZ[p], RZ[p+1], 0);

        vtkIdType tri[3];
        for (int i = 0; i < numTri; i++)
        {
            tri[0] = conn[i*3+0];
            tri[1] = conn[i*3+1];
            tri[2] = conn[i*3+2];
            grid->InsertNextCell(VTK_TRIANGLE, 3, tri);
        }
    }
    //It's a 3D mesh
    else
    {
        int numPlanes = numPhi+1;
        
        double dPhi = 2.0*M_PI/(double)(numPlanes-1);
        points->SetNumberOfPoints(numNodes * numPlanes);
        for (int i = 0; i < numPlanes; i++)
        {
            double phi = - (double)i * dPhi;
            int p = 0;
            for (int j = 0; j < numNodes; j++, p+=2)
                points->SetPoint(j+i*numNodes, RZ[p], phi, RZ[p+1]);
        }

        vtkIdType wedge[6];
        for (int i = 0; i < numPlanes-1; i++)
        {
            int off = i*(numNodes);            
            for (int p = 0; p < numTri; p++)
            {
                int off = i*(numNodes);
                int p0 = conn[p*3+0];
                int p1 = conn[p*3+1];
                int p2 = conn[p*3+2];
                
                wedge[0] = p0 + off;
                wedge[1] = p1 + off;
                wedge[2] = p2 + off;
                //grid->InsertNextCell(VTK_TRIANGLE, 3, wedge);
                
                off = (i+1)*(numNodes);
                wedge[3] = p0 + off;
                wedge[4] = p1 + off;
                wedge[5] = p2 + off;
                grid->InsertNextCell(VTK_WEDGE, 6, wedge);
            }
        }
    }

    delete [] RZ;
    delete [] conn;
    H5Fclose(file);
    
    return grid;
}


// ****************************************************************************
//  Method: avtXGCFileFormat::GetVar
//
//  Purpose:
//      Gets a scalar variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      varname    The name of the variable requested.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Jun 7 11:27:14 PDT 2016
//
// ****************************************************************************

vtkDataArray *
avtXGCFileFormat::GetVar(const char *varname)
{
    hid_t file;
    if ((file = H5Fopen(GetFilename(), H5F_ACC_RDONLY, H5P_DEFAULT)) < 0)
        EXCEPTION1(InvalidFilesException, GetFilename());

    hsize_t dims[10];
    hid_t grp = H5Gopen(file, "/", H5P_DEFAULT);
    hid_t vD = H5Dopen(grp, varname, H5P_DEFAULT);
    hid_t vS = H5Dget_space(vD);
    int n = H5Sget_simple_extent_dims(vS, dims, NULL);

    int sz = (n==1 ? dims[0] : (dims[0]*dims[1] + dims[0]));

    vtkFloatArray *arr = vtkFloatArray::New();
    arr->SetNumberOfTuples(sz);

    if (n == 1)
        H5Dread(vD, H5T_NATIVE_FLOAT, H5S_ALL, vS, H5P_DEFAULT, arr->GetVoidPointer(0));
    else
    {
        std::vector<float> tmp(sz);
        H5Dread(vD, H5T_NATIVE_FLOAT, H5S_ALL, vS, H5P_DEFAULT, &tmp[0]);

        //Need to transpose the array...
        int d0 = dims[1], d1 = dims[0];
        for (int r = 0; r < d0; r++)
            for (int c = 0; c < d1; c++)
                arr->SetTuple1(d1*(d0-1-r)+c, tmp[c*d0 + r]);

        //copy the first plane to last..
        int off = numNodes*numPhi;
        for (int i = 0; i < numNodes; i++)
            arr->SetTuple1(i+off, arr->GetTuple1(i));
    }
    
    H5Dclose(vD);
    H5Gclose(grp);
    H5Fclose(file);

    return arr;
}


// ****************************************************************************
//  Method: avtXGCFileFormat::GetVectorVar
//
//  Purpose:
//      Gets a vector variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      varname    The name of the variable requested.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Jun 7 11:27:14 PDT 2016
//
// ****************************************************************************

vtkDataArray *
avtXGCFileFormat::GetVectorVar(const char *varname)
{
    return NULL;
}


// ****************************************************************************
//  Method: avtXGCFileFormat::Initialize
//
//  Purpose:
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Tue Jun 7 11:27:14 PDT 2016
//
// ****************************************************************************

void
avtXGCFileFormat::Initialize()
{
    const char *mName = "avtGTCFileFormat::Initialize: ";

    if (initialized)
        return;

    // Init HDF5 and turn off error message printing.
    H5open();
    H5Eset_auto( H5E_DEFAULT, NULL, NULL );

    // Check for a valid XGC file
    if( H5Fis_hdf5( GetFilename() ) < 0 )
        EXCEPTION1( InvalidFilesException, GetFilename() );

    hid_t fileHandle;
    if ((fileHandle = H5Fopen(GetFilename(), H5F_ACC_RDONLY, H5P_DEFAULT)) < 0)
        EXCEPTION1( InvalidFilesException, GetFilename() );

    hsize_t dims[1];
    
    hid_t nphiD = H5Dopen(fileHandle, "nphi", H5P_DEFAULT);
    if(nphiD < 0)
    {
      H5Fclose(fileHandle);
      EXCEPTION1( InvalidFilesException, GetFilename() );
    }
    else
    {
      hid_t nphiS = H5Dget_space(nphiD);
      H5Sget_simple_extent_dims(nphiS, dims, NULL);
      H5Dread(nphiD, H5T_NATIVE_INT, H5S_ALL, nphiS, H5P_DEFAULT, &numPhi);
      H5Dclose(nphiD);
      H5Sclose(nphiS);
    }

    
    hid_t nnodeD = H5Dopen(fileHandle, "nnode", H5P_DEFAULT);
    if(nnodeD < 0)
    {
      H5Fclose(fileHandle);
      EXCEPTION1( InvalidFilesException, GetFilename() );
    }
    else
    {
      hid_t nnodeS = H5Dget_space(nnodeD);
      H5Sget_simple_extent_dims(nnodeS, dims, NULL);
      H5Dread(nnodeD, H5T_NATIVE_INT, H5S_ALL, nnodeS, H5P_DEFAULT, &numNodes);
      H5Dclose(nnodeD);
      H5Sclose(nnodeS);
    }

    H5Fclose(fileHandle);

    // At this point consider the file to truly be a XGC file. If
    // some other file exception will be thrown.

    if( numPhi <= 0 )
    {
      debug4 << mName << "Could not determine number of phi slices" << std::endl;
      
      H5Fclose(fileHandle);
      EXCEPTION1( InvalidVariableException, "XGC Dataset Extents - Dataset 'nPhiD' has an invalid value");
    }

    if(numNodes <= 0)
    {
      debug4 << mName << "Could not determine number of nodes" << std::endl;
      H5Fclose(fileHandle);
      EXCEPTION1( InvalidVariableException, "XGC Dataset Extents - Dataset 'nnodeD' has an invalid value");
    }
    
    initialized = true;
}
